#include "nm_workthread.h"

kworkthread::kworkthread() : m_Stop(FALSE), m_Lookaside(sizeof(WORK_ITEM))
{
	RtlZeroMemory(&m_ClientId, sizeof(CLIENT_ID));
	InitializeListHead(&m_ListHeader);
	KeInitializeEvent(&m_InsertEvent, SynchronizationEvent, FALSE);
	KeInitializeSpinLock(&m_Lock);
	PsCreateSystemThread(&m_ThreadHandle, THREAD_ALL_ACCESS, NULL, NULL, &m_ClientId, Run, this);
}

void kworkthread::Release()
{
	m_Stop = TRUE;
	KeSetEvent(&m_InsertEvent, 0, FALSE);
	KeWaitForSingleObject(m_ThreadHandle, Executive, KernelMode, FALSE, NULL);
	ZwClose(m_ThreadHandle);
}


void kworkthread::Run(PVOID Context)
{
	kworkthread *pthis = (kworkthread *)Context;
	LARGE_INTEGER timeout;
	PWORK_ITEM WorkItem;
	NTSTATUS ns;
	timeout.QuadPart = -10000 * 10 * 1000;

	while (TRUE)
	{
		ns = KeWaitForSingleObject((PVOID)&pthis->m_InsertEvent, Executive, KernelMode, FALSE, &timeout);
		if (ns == STATUS_TIMEOUT) {

			KdPrint(("kworkthread timeout\n"));
		}

		if (pthis->m_Stop) { 
			break;
		}

		while ((WorkItem = 
			CONTAINING_RECORD( 
			ExInterlockedRemoveHeadList(&pthis->m_ListHeader, &pthis->m_Lock), 
			WORK_ITEM, 
			Next )) != NULL)
		{
			if (WorkItem->StartAddress != NULL) {
				WorkItem->StartAddress(WorkItem->Context);
			}

			KeSetEvent( &WorkItem->Event, 0, FALSE );
			pthis->m_Lookaside.Free(WorkItem);
		}
		
		if (pthis->m_Stop) {
			break;
		}
	}
}

void kworkthread::InsertItem(WORKITEMPROC Address, PVOID Context, PRKEVENT *Event)
{
	PWORK_ITEM WorkItem;

	if (!MmIsAddressValid(Address)) {
		return;
	}

	WorkItem = (PWORK_ITEM)m_Lookaside.Allocate();
	if (WorkItem == NULL) return;

	WorkItem->Next.Flink = NULL;
	WorkItem->Next.Blink = NULL;
	WorkItem->StartAddress = Address;
	WorkItem->Context = Context;
	if (Event != NULL) {
		*Event = &WorkItem->Event;
	}
	KeInitializeEvent(&WorkItem->Event, SynchronizationEvent, FALSE);
	ExInterlockedInsertTailList(&m_ListHeader, &WorkItem->Next, &m_Lock);
	KeSetEvent(&m_InsertEvent, 0, FALSE);
}